#include <vector>
using std::vector;
using std::pair;

#include "opencv\cv.h"
#include "opencv\cxcore.h"
#include "opencv\highgui.h"
#include "opencv\ml.h"

#define BACKGROUND_VALUE            0
#define FOREGROUND_VALUE          255

#define BAYESIAN_NUMBER_NEAREST   200
#define BAYESIAN_SIGMA            8.f
#define BAYESIAN_SIGMA_C          5.f
#define BAYESIAN_MAX_CLUS           3

#define FLASH_NUMBER_NEAREST      200
#define FLASH_SIGMA               8.f
#define FLASH_SIGMA_I_SQUARE     32.f
#define FLASH_SIGMA_IP_SQUARE    32.f
#define FLASH_MAX_CLUS              3

class Matting
{
public:
	Matting(){ trimap = alphamap = NULL; }
	virtual double Solve() = 0;
	IplImage *trimap;
	IplImage *alphamap;
};


//Implementation of	"A Bayesian Approach to Digital Matting"
class BayesianMatting : public Matting
{
public:
	/* ==================================================================================
		Constructors/ Destructors. 
	   ================================================================================== */

	//the input is the color image and the trimap image
	//the format is 3 channels + uchar and 1 channel + uchar respectively
	BayesianMatting( IplImage* cImg, IplImage* trimap );
	~BayesianMatting();
	
	void Initialize();

	//set parameter
	void SetParameter( int N = BAYESIAN_NUMBER_NEAREST, float sigma_ = BAYESIAN_SIGMA, float sigma_c = BAYESIAN_SIGMA_C );

	//solve the matting problem
	double Solve();

	

private:	
	/* ==================================================================================
		Internal functions.
	   ================================================================================== */	
	//get the extreme outer contours of an image
	void GetContour( IplImage* img, vector<CvPoint> &contour );

	//initialize the alpha of one point using the mean of neighors and save the result in alphamap
	void InitializeAlpha( int r, int c, const IplImage* unSolvedMask );

	//used for clustering according the equation in paper "Color Quantization of Images"
	void CalculateNonNormalizeCov( const IplImage *cImg, const vector<pair<CvPoint,float>> &clus_set, CvMat* mean, CvMat* cov );

	// calculate mean and cov of the given clus_set
	void CalculateMeanCov( const IplImage *cImg, const vector<pair<CvPoint,float>> &clus_set, CvMat* mean, CvMat* cov );
	
	// calculate weight, mean and cov of the given clus_set
	void CalculateWeightMeanCov( const IplImage *cImg, const vector<pair<CvPoint,float>> &clus_set, float &weight, CvMat* mean, CvMat* cov );

	//get the foreground and backgroud gmm model at a given pixel
	void GetGMMModel( int r, int c, vector<float> &fg_weight, const vector<CvMat*> fg_mean, const vector<CvMat*> inv_fg_cov, vector<float> &bg_weight, const vector<CvMat*> bg_mean, const vector<CvMat*> inv_bg_cov );

	//collect the foreground/background sample set from the contour, called by GetGMMModel
	void CollectSampleSet( int r, int c, vector<pair<CvPoint, float>> &fg_set, vector<pair<CvPoint, float>> &bg_set );

	//solve Eq. (9) at pixel (r,c) according to the alpha in alphamap and save the result in fgImg and bgImg
	void SolveBF( int r, int c, CvMat *fg_mean, CvMat *inv_fg_cov, CvMat *bg_mean, CvMat *inv_bg_cov );

	//solve Eq. (10) at pixel (r,c) according to the foreground and background color in fgImg and bgImg, and save the result in alphamap
	inline void SolveAlpha( int r, int c );

	//compute total likelihood
	float computeLikelihood(int r, int c, CvMat *fg_mean, CvMat *inv_fg_cov, CvMat *bg_mean, CvMat *inv_bg_cov);

	float computeLikelihood(int r, int c, float fg_weight, CvMat *fg_mean, CvMat *inv_fg_cov, float bg_weight, CvMat *bg_mean, CvMat *inv_bg_cov);


	/* ==================================================================================
		Offline Variables.
	   ================================================================================== */
	int   nearest;	
	float sigma;
	float sigmac;
	
	/* ==================================================================================
		Online Variables.
	   ================================================================================== */
	
	IplImage *colorImg, *fgImg, *bgImg;
	IplImage *bgmask, *fgmask, *unmask, *unsolvedmask;
};

//Implementation of	"Flash Matting"
class FlashMatting : public Matting
{
public:
	/* ==================================================================================
		Constructors/ Destructors. 
	   ================================================================================== */

	//the input is the ambient image (no flash), flash image, and the trimap image is optional
	//the format is 3 channels + uchar 3 channels + uchar and 1 channel + uchar respectively
	FlashMatting( IplImage* ambientImg, IplImage* flashImg, IplImage* trimap = NULL );
	~FlashMatting();

	void Initialize();

	//set parameter
	void SetParameter( int N = FLASH_NUMBER_NEAREST, float sigma_ = FLASH_SIGMA, float sigma_i_square = FLASH_SIGMA_I_SQUARE, float sigma_ip_square = FLASH_SIGMA_IP_SQUARE );
	
	//automatically generate the trimap from flash/no-flash image pair
	static void GenerateTrimap( IplImage *FlashOnlyImg, IplImage *&Trimap, float level = 1.f, float ratio = 0.6f );

	//solve the matting problem
	double Solve();
private:
	/* ==================================================================================
		Internal functions.
	   ================================================================================== */	
	//get the extreme outer contours of an image
	void GetContour( IplImage* img, vector<CvPoint> &contour );
	
	//get the model at a given pixel
	void GetGMMModel( int r, int c, const IplImage *cimg, const vector<pair<CvPoint,float>> &point_set, vector<float> &weight, const vector<CvMat*> mean, const vector<CvMat*> inv_cov );

	//used for clustering according the equation in paper "Color Quantization of Images"
	void CalculateNonNormalizeCov( const IplImage *cImg, const vector<pair<CvPoint,float>> &clus_set, CvMat* mean, CvMat* cov );
	
	// calculate weight, mean and cov of the given clus_set
	void CalculateWeightMeanCov( const IplImage *cImg, const vector<pair<CvPoint,float>> &clus_set, float &weight, CvMat* mean, CvMat* cov );

	//collect the foreground/background sample set from the contour, called by GetGMMModel
	void CollectSampleSet( int r, int c, vector<pair<CvPoint, float>> &point_set, vector<pair<CvPoint, float>> &bg_set );
	
	//initialize the alpha of one point using the mean of neighors and save the result in alphamap
	void InitializeAlpha( int r, int c, const IplImage* unSolvedMask );

	//solve Eq. (12) at pixel (r,c) according to the foreground, background color, and flash-only-fg color, and save the result in alphamap
	inline void SolveAlpha( int r, int c );

	//solve Eq. (13) at pixel (r,c) according to the alpha in alphamap and save the result in fgImg, flashFgimg, and bgImg
	void SolveBFF_( int r, int c, CvMat *fg_mean, CvMat *inv_fg_cov, CvMat *bg_mean, CvMat *inv_bg_cov, CvMat *flash_fg_mean, CvMat *inv_flash_fg_cov );
	
	//compute total likelihood
	float computeLikelihood(int r, int c, CvMat *fg_mean, CvMat *inv_fg_cov, CvMat *bg_mean, CvMat *inv_bg_cov, CvMat *flash_fg_mean, CvMat *inv_flash_fg_cov );
	float computeLikelihood( int r, int c, float fg_weight, CvMat *fg_mean, CvMat *inv_fg_cov, float bg_weight, CvMat *bg_mean, CvMat *inv_bg_cov, float flash_fg_weight, CvMat *flash_fg_mean, CvMat *inv_flash_fg_cov );
	/* ==================================================================================
		Offline Variables.
	   ================================================================================== */
	int   nearest;	
	float sigma;
	float sigma_i_square, sigma_ip_square;
	
	/* ==================================================================================
		Online Variables.
	   ================================================================================== */
	
	IplImage *ambientImg, *flashImg, *flashOnlyImg;
	IplImage *fgImg, *flashFgImg, *bgImg;
	IplImage *bgmask, *fgmask, *unmask, *unsolvedmask;
};
